MODULE qcim_main

   USE qcim_global_types
   USE qcim_global_consts
   USE md1_global_types
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: qcim_initial,                  &
             qcim_final,                    &
             qcim_get_Q_matrix

   TYPE(md1_FE_grid), SAVE :: grid
   INTEGER(INTK), SAVE :: use_images
   REAL(REALK) :: pshift = 0d0

CONTAINS

!-------------------------------------------------------------------------------

   SUBROUTINE qcim_initial(nbas)

      USE qcim_geometry,      ONLY: qcim_define_geometry
      USE qcim_scheme_module, ONLY: qcim_init_scheme, qcim_get_scheme
      USE fem_main,           ONLY: fem_initial
      USE fmm_main,           ONLY: fmm_initial
      USE md1_main,           ONLY: md1_initial, md1_write_multipole_integrals

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: nbas

      INTEGER(INTK), PARAMETER :: fem_order = 3
      INTEGER(INTK), PARAMETER :: multipole_order = 8
      TYPE(qcim_scheme), POINTER :: scheme
      INTEGER(INTK) :: ngrid(3), error, i
      REAL(REALK)   :: origin(3), lambda(3)

      OPEN (UNIT=IO1,FILE='qcim.out',STATUS='REPLACE',    &
            ACCESS='SEQUENTIAL',FORM='FORMATTED')
      REWIND (IO1)

      WRITE(IO1,*)
      WRITE(IO1,*) "-----------------------------------------"
      WRITE(IO1,*) "|         QCIM initialization           |"
      WRITE(IO1,*) "-----------------------------------------"
      WRITE(IO1,*)

      ! Define QCIM geometry and fixed boundary potential
      CALL qcim_define_geometry(fem_order,ngrid,origin,lambda,   &
                                use_images,pshift)

      ! Initialize runtime database
      CALL qcim_init_scheme(0,ngrid)
      CALL qcim_get_scheme(scheme)

      ! Initialize FEM geometry and grid
      CALL fem_initial('XXX',ngrid,lambda)
      grid%order = fem_order
      grid%h(:) = lambda(:)
      grid%L(:) = lambda(:) * DBLE(grid%order)
      grid%n(:) = ngrid(:)
      grid%lo(:) = origin(:)
      grid%hi(:) = origin(:) + DBLE(grid%n(:)-1) * grid%h(:)
      WRITE(IO1,*) "FEM volume:"
      WRITE(*,*) "FEM volume:"
      DO i = 1,3
         WRITE(IO1,'(I3,3F15.5)') i, grid%lo(i), grid%hi(i), grid%h(i)
         WRITE(*,'(I3,3F15.5)') i, grid%lo(i), grid%hi(i), grid%h(i)
      END DO

      ! Grid is now fixed, so allocate grid object for all SCF iterations
      ALLOCATE (grid%f(ngrid(1),ngrid(2),ngrid(3)),STAT=error)
      IF (error /= 0) STOP 'FEM grid allocation failure'
      grid%f(:,:,:) = zero

      ! Initialize MD1 and FMM libraries
      CALL md1_initial('basis.qcim',nbas)
      CALL md1_write_multipole_integrals(multipole_order)
      CALL fmm_initial('QCIM_FMM',multipole_order)
      !CALL fmm_initial('MD4_FMM',multipole_order)

   END SUBROUTINE qcim_initial

!-------------------------------------------------------------------------------

   SUBROUTINE qcim_final

      USE fmm_main, ONLY: fmm_final
      USE fem_main, ONLY: fem_final
      USE md1_main, ONLY: md1_final
      IMPLICIT NONE
      CALL qcim_compute_nuclear_correction
      CALL fmm_final
      CALL fem_final
      CALL md1_final
      IF (ASSOCIATED (grid%f)) DEALLOCATE (grid%f)
      NULLIFY (grid%f)
      CLOSE (UNIT=IO1,STATUS='KEEP')

   END SUBROUTINE qcim_final

!-------------------------------------------------------------------------------

   SUBROUTINE qcim_set_laplace_BC(scheme,dens)

      USE fmm_main, ONLY: fmm_get_multipole_potential

      IMPLICIT NONE
      TYPE(qcim_scheme), INTENT(IN) :: scheme
      REAL(REALK),       INTENT(IN) :: dens(:,:)

      REAL(REALK) :: fixed_pot(scheme%nbpoints)
      REAL(REALK) :: fmm_pot(1,scheme%nbpoints)
      REAL(REALK) :: coor(3,scheme%nbpoints)
      INTEGER(INTK) :: ixyz(3,scheme%nbpoints)
      INTEGER(INTK) :: ix,iy,iz, imax=0,jmax=0,kmax=0
      INTEGER(INTK) :: mode = 0

      INTEGER(INTK) :: i, nbpts

      ! Read in fixed QCIM boundary potential
      WRITE(IO1,*) "Reading in fixed BC..."
      OPEN (UNIT=IO2,FILE='bpot.qcim',STATUS='OLD',    &
            ACTION='READ',FORM='UNFORMATTED')
      REWIND (IO2)
      READ (IO2) nbpts 
      IF (nbpts /= scheme%nbpoints) STOP 'fixme'
      DO i = 1, nbpts
         READ (IO2) ix,iy,iz, fixed_pot(i)
         imax = MAX(imax,ix)
         jmax = MAX(jmax,iy)
         kmax = MAX(kmax,iz)
!LizaFIXME
!WRITE(IO1,*) "BC:", ix,iy,iz, fixed_pot(i)
         ixyz(:,i) = (/ix,iy,iz/)
         coor(:,i) = grid%lo(:) + DBLE((/ix-1,iy-1,iz-1/)) * grid%h(:)
      END DO
      CLOSE(UNIT=IO2,STATUS='KEEP')

      ! Get molecular (nuclear + electron) boundary potential via FMM
      WRITE(IO1,*) "Getting molecular FMM free-space BC..."
      CALL fmm_get_multipole_potential(mode,nbpts,coor,dens,fmm_pot)

      ! Open FEM solver data file and write boundary potential 
      ! due to "external" charges only (i.e. fixed - molecular)
      WRITE(IO1,*) "Writing QCIM BC..."
      OPEN(UNIT=IO2, FILE='bpot.fem', STATUS='REPLACE',  &
           ACCESS='SEQUENTIAL', FORM='UNFORMATTED')
      REWIND(IO2)
      WRITE (IO2) nbpts 

      WRITE(IO1,*) "------------------------------------------"
      WRITE(IO1,*) "Free-space potential:"
      WRITE(IO1,*) "------------------------------------------"

      DO i = 1, nbpts
         IF (use_images == 1) THEN
            ! Fix total potential at this point as specified "fixed_pot"
            WRITE(IO2) ixyz(1:3,i), fixed_pot(i) - fmm_pot(1,i) + pshift
         ELSE IF (use_images == 0) THEN
            ! Mimic static external field:
            WRITE(IO2) ixyz(1:3,i), fixed_pot(i) + pshift
         ELSE IF (use_images == 2) THEN
            ! Mimic static external field (in GHOST mode)
            WRITE(IO2) ixyz(1:3,i), fixed_pot(i) + pshift
         ELSE
            STOP 'use_images error!'
         END IF

         ! Print out some of them at centers of plates:
         ! Bottom plate
         IF ( (ixyz(1,i) == (imax+1)/2) .AND.  &
              (ixyz(2,i) == (jmax+1)/2) .AND.  &
              (ixyz(3,i) == 1)) THEN
            WRITE(IO1,'(3I,3F12.6,E20.10)'), ixyz(:,i), coor(:,i), fmm_pot(1,i)
         END IF
         ! Top plate
         IF ( (ixyz(1,i) == (imax+1)/2) .AND.  &
              (ixyz(2,i) == (jmax+1)/2) .AND.  &
              (ixyz(3,i) == kmax)) THEN
            WRITE(IO1,'(3I,3F12.6,E20.10)'), ixyz(:,i), coor(:,i), fmm_pot(1,i)
         END IF
         ! Front plate
         IF ( (ixyz(1,i) == (imax+1)/2) .AND.  &
              (ixyz(2,i) == 1) .AND.  &
              (ixyz(3,i) == (kmax+1)/2)) THEN
            WRITE(IO1,'(3I,3F12.6,E20.10)'), ixyz(:,i), coor(:,i), fmm_pot(1,i)
         END IF
         ! Back plate
         IF ( (ixyz(1,i) == (imax+1)/2) .AND.  &
              (ixyz(2,i) == jmax) .AND.  &
              (ixyz(3,i) == (kmax+1)/2)) THEN
            WRITE(IO1,'(3I,3F12.6,E20.10)'), ixyz(:,i), coor(:,i), fmm_pot(1,i)
         END IF
         ! Side plate x-min
         IF ( (ixyz(1,i) == 1) .AND.  &
              (ixyz(2,i) == (jmax+1)/2) .AND.  &
              (ixyz(3,i) == (kmax+1)/2)) THEN
            WRITE(IO1,'(3I,3F12.6,E20.10)'), ixyz(:,i), coor(:,i), fmm_pot(1,i)
         END IF
         ! Side plate x-max
         IF ( (ixyz(1,i) == imax) .AND.  &
              (ixyz(2,i) == (jmax+1)/2) .AND.  &
              (ixyz(3,i) == (kmax+1)/2)) THEN
            WRITE(IO1,'(3I,3F12.6,E20.10)'), ixyz(:,i), coor(:,i), fmm_pot(1,i)
         END IF

      END DO

      WRITE(IO1,*) "------------------------------------------"
      WRITE(IO1,*) "Fixed potential:"
      WRITE(IO1,*) "------------------------------------------"
      WRITE(IO1,*) "(1,1,1):"
      WRITE(IO1,'(3F12.6,E20.10)'), coor(:,1), fixed_pot(1)
      WRITE(IO1,*) "(N,N,N):"
      WRITE(IO1,'(3F12.6,E20.10)'), coor(:,nbpts), fixed_pot(nbpts)
      WRITE(IO1,*) "------------------------------------------"
      WRITE(IO1,*) "Constant input z shift", pshift
      WRITE(IO1,*) "------------------------------------------"

      CLOSE(UNIT=IO2, STATUS='KEEP')

   END SUBROUTINE qcim_set_laplace_BC

!-------------------------------------------------------------------------------

   SUBROUTINE qcim_get_Q_matrix(nbas,dens,Q_matrix)

      USE qcim_scheme_module, ONLY: qcim_get_scheme
      USE fem_main,           ONLY: fem_solve_laplace_equation
      USE md1_main,           ONLY: md1_get_AO_FE_matrix_elements

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: nbas
      REAL(REALK),   INTENT(IN)  :: dens(nbas,nbas)
      REAL(REALK),   INTENT(OUT) :: Q_matrix(nbas,nbas)

      TYPE(qcim_scheme), POINTER :: scheme
      REAL(REALK) :: foo, S(nbas,nbas)
      INTEGER(INTK) :: i,j,k, P,Q
      REAL(REALK) :: xyz(3)
      
      WRITE(IO1,*) "------------------------------------------"
      WRITE(IO1,*) "Building Q matrix... Hi Liza!"
      CALL qcim_get_scheme(scheme)

      ! Reset the solver boundary conditions based on this density
      WRITE(IO1,*) "Updating Laplace BC..."
      print *, "Updating Laplace BC..."
      CALL qcim_set_laplace_BC(scheme,dens)

      ! Solve Laplace and overlap solution with AOs
      WRITE(IO1,*) "Solving Laplace..."
      print *, "Solving Laplace..."
      CALL fem_solve_laplace_equation('bpot.fem',grid%f)

!LizaFIXME
!DO k = 1, grid%n(3)
!DO j = 1, grid%n(2)
!DO i = 1, grid%n(1)
!   xyz(:) = grid%lo(:) + DBLE((/i-1,j-1,k-1/)) * grid%h(:)
!   WRITE(IO1,*) "Veff(r)", xyz(:), grid%f(i,j,k) 
!END DO
!END DO
!END DO

      WRITE(IO1,*) "Building AO/FE overlap..."
      print *, "Building AO/FE overlap..."
      CALL md1_get_AO_FE_matrix_elements(grid,Q_matrix)
!Q_matrix(:,:) = 0.01*Q_matrix(:,:) 
      print *, "done!"

      CALL qcim_dump_grid_AO_overlap(nbas,Q_matrix)

!----------- diagnostics -----------
      ! Check format of AO density matrix against the test
      ! N = Sum_{ab} S_ab D_ab
      ! where N = total number of electrons
      OPEN (UNIT=IO2,FILE='overlap_matrix.md1',STATUS='OLD',  &
            ACTION='READ',FORM='FORMATTED')
      REWIND (IO2)
      S(:,:) = 0d0
      DO i = 1, nbas*nbas
         READ (IO2,'(3I8,F25.15)') k, P,Q, foo
         IF (P == -1) EXIT
         S(p,q) = foo
         S(q,p) = foo
!         WRITE(IO1,*) "density(p,q) =", p,q, dens(p,q)
!         WRITE(IO1,*) "S(p,q)       =", p,q, S(p,q)
      END DO
      CLOSE (UNIT=IO2)
      ! Note we only have tri-loop, but off-diagonal Dab
      ! elements have a factor of 2 inside them!!
      foo = 0d0
      DO p = 1, nbas
         DO q = p, nbas
            foo = foo + S(p,q)*dens(p,q)
         END DO
      END DO
      WRITE(IO1,*) "Density matrix check: N =", foo

      ! Electronic-V(eff) interaction check
      foo = 0d0
      DO i = 1, nbas
         DO j = i, nbas
            foo = foo + Q_matrix(j,i)*dens(j,i)
         END DO
      END DO
      WRITE(IO1,*) "check: Sum_{ab} Q-matrix(a,b) * Dens(a,b) =", -foo
!----------- diagnostics -----------

      ! Apply input flag to NOT add Q-matrix to Fock matrix (Ghost mode)
      IF (use_images == 2) THEN
         Q_matrix(:,:) = zero
      END IF

   END SUBROUTINE qcim_get_Q_matrix

!-------------------------------------------------------------------------------

   SUBROUTINE qcim_dump_grid_AO_overlap(nbas,Q_matrix)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: nbas
      REAL(REALK),   INTENT(IN) :: Q_matrix(nbas,nbas)
      INTEGER(INTK) :: p,q

      OPEN (UNIT=IO2, FILE='Q_matrix.qcim', STATUS='REPLACE',  &
            ACTION='WRITE', FORM='FORMATTED')
      REWIND (IO2)

      WRITE(IO1,*) "Dumping AO/FE field overlap matrix..."
      DO p = 1, nbas
         DO q = 1, p
            IF (ABS(Q_matrix(p,q)) > 1d-13)  &
               WRITE(IO2,'(2I8,F25.15)') p,q, Q_matrix(p,q)
         END DO
      END DO
      CLOSE (UNIT=IO2, STATUS='KEEP')

   END SUBROUTINE qcim_dump_grid_AO_overlap

!-------------------------------------------------------------------------------

   SUBROUTINE qcim_compute_nuclear_correction

      USE md1_main, ONLY: md1_get_FE_charge_overlap

      IMPLICIT NONE
      INTEGER(INTK) :: i, nnuc, Q

      REAL(REALK), ALLOCATABLE :: centres(:,:)
      REAL(REALK), ALLOCATABLE :: charges(:)
      REAL(REALK), ALLOCATABLE :: energies(:,:)

      OPEN (UNIT=IO2,FILE='nuclei.qcim',STATUS='OLD',  &
            ACTION='READ',FORM='FORMATTED')
      REWIND (IO2)
      READ (IO2,'(I8)') nnuc

      ALLOCATE (centres(3,nnuc))
      ALLOCATE (charges(nnuc))
      ALLOCATE (energies(nnuc,nnuc))

      DO i = 1, nnuc
         READ (IO2,'(F25.15)') centres(1,i)
         READ (IO2,'(F25.15)') centres(2,i)
         READ (IO2,'(F25.15)') centres(3,i)
         READ (IO2,'(I5)') Q
         charges(i) = DBLE(Q)
      END DO
      CLOSE(UNIT=IO2, STATUS='KEEP')

      ! Compute interaction energy for each nucleus in the QCIM potential
      energies = zero
      CALL md1_get_FE_charge_overlap(nnuc,charges,centres,grid,energies)

      WRITE(IO1,*) "Nuclear QCIM correction:", SUM(energies)
      print *, "Nuclear QCIM correction:", SUM(energies)
DO i = 1, nnuc
   print *, i, charges(i), energies(i,i)
END DO

      DEALLOCATE (centres)
      DEALLOCATE (charges)
      DEALLOCATE (energies)

   END SUBROUTINE qcim_compute_nuclear_correction

!-------------------------------------------------------------------------------

END MODULE qcim_main

